import os
import numpy as np
import matplotlib
import matplotlib.pyplot as plt
import pdb

from config import if_compare_plot_slacks

matplotlib.use('TkAgg')

for b in ['usbf_device']:
    data = np.load('./checkpoints/slacks/{}.npz'.format(b))
    se, sl, se_truth, sl_truth = [data['arr_{}'.format(i)].flatten() for i in range(4)]

    data_original = np.load('./checkpoints/original_slacks/{}.npz'.format(b))#文章中的slacks
    se_original, sl_original, se_truth_original, sl_truth_original = [data_original['arr_{}'.format(i)].flatten() for i in range(4)]

    def plot_slack(ax, se_truth, se, title, color_):
        # pdb.set_trace()  #程序运行到这里就会暂停
        ax.set_title(title)
        ax.scatter(se_truth, se, s=1, color=color_)
        ax.set_xlabel('Truth')
        ax.set_ylabel('Predicted')
        # line = np.polyfit(se_truth, se, 1)
        # y1 = line[0] * se_truth + line[1]
        # ax.plot(se_truth, y1, 'r-')
        ax.plot(se_truth, se_truth, 'k-')
        ax.plot(se_truth, se, 'rx')

    fig1, ax1 = plt.subplots(1, 1)
    fig2, ax2 = plt.subplots(1, 1)
    # fig.tight_layout()
    #plt.subplots_adjust(left=0.125, bottom=0.2, right=0.9, top=0.9, wspace=0.4, hspace=0.2)
    #Hold slack=data arrival time -data required time
    #Setup slack=data required time -data arrival time
    plot_slack(ax1, se_truth, se, 'Hold slacks', 'red')
    ax1.xaxis.set_ticks([-5, -2.5, 0, 2.5, 5])
    ax1.yaxis.set_ticks([-5, -2.5, 0, 2.5, 5])
    plot_slack(ax2, sl_truth, sl, 'Setup slacks', 'red')
    ax2.xaxis.set_ticks([-10, -5, 0, 5, 10])
    ax2.yaxis.set_ticks([-10, -5, 0, 5, 10])

    if if_compare_plot_slacks:
        plot_slack(ax1, se_truth_original, se_original, 'Hold slacks', None)
        plot_slack(ax2, sl_truth_original, sl_original, 'Setup slacks', None)

    fig1.canvas.manager.set_window_title(b)
    fig2.canvas.manager.set_window_title(b)
    fig1.savefig(f'./checkpoints/slacks_fig/{b}_hold.png')
    fig2.savefig(f'./checkpoints/slacks_fig/{b}_setup.png')
    plt.show()
